#include "avl.h"



avl_t* creat_avl()
{
    avl_t* tree = malloc(sizeof(avl_t));
    if (tree == NULL) {
        return NULL;
    }
    memset(tree, 0, sizeof(avl_t));
    
    return tree;
}

int destroy_avl(avl_t* tree)
{
    avl_node_t* stack[10000];
    int top = -1;
    avl_node_t *cur;

    stack[++top] = tree->root;
    cur = tree->root->left;

    while (cur != NULL || top != -1 ) {
        while (cur != NULL) {
            stack[++top] = cur;
            cur = cur->left;
        }
        
        cur = stack[top]->right;

        free(stack[top--]);
        tree->node_count--;
    }

    free(tree);
    return 0;
}



/**
 * @brief 计算节点是否是叶子节点
 * 
 * @param node 
 * @return 是返回1 不是返回0
 */
int is_leaf_node(avl_node_t* node)
{
    if (node == NULL) {
        return 0;
    }

    return node->left == NULL && node->right == NULL;
}

/**
 * @brief 创建节点
 * 
 * @param key: 键值
 * @return 新节点的指针
 */
avl_node_t* creat_node(int key)
{
    avl_node_t* new_node = malloc(sizeof(avl_node_t));
    if (new_node == NULL) {
        return NULL;
    }
    memset(new_node, 0, sizeof(avl_node_t));

    new_node->height = 1;
    new_node->val = key;

    return new_node;
}

/*! 寻找目标节点，找到返回节点，没找到返回NULL */
/**
 * @brief   寻找目标节点，找到返回节点，没找到返回NULL
 * 
 * @param   root: 从此节点开始找
 * @param   key: 键值
 * @return  找到了返回目标节点，没找到返回NULL
 */
avl_node_t* search_node(avl_node_t* root, int key)
{
    avl_node_t* target = root;

    if(root == NULL) {
        return NULL;
    }

    while (target != NULL) {
        if (key < target->val)
        {
            target = target->left;
        } else if (key > target->val) {
            target = target->right;
        } else {
            return target;
        }
    }

    /* 没有目标节点 */
    return NULL;
}

/**
 * @brief 给树上插入一个新节点
 * 
 * @param tree 树
 * @param key 键值
 * @return 返回树本身
 */
avl_t* insert_node(avl_t* tree, int key)
{
    avl_node_t* cur = NULL;
    avl_node_t* new_node = NULL;


    new_node = creat_node(key);
    if (new_node == NULL) {
        return tree;
    }

    // 如果是空树，则创建根结点
    if (tree->root == NULL) {
        tree->root = new_node;
        tree->node_count++;
        return tree;
    }
    
    cur = tree->root;
    new_node->parent = cur;
    while (cur != NULL) {
        new_node->parent = cur;

        if (key < cur->val) {
            if (cur->left == NULL) {
                cur->left = new_node;
                break;
            } else {
                cur = cur->left;
            }
        } else if (key > cur->val) {
            if (cur->right == NULL) {
                cur->right = new_node;
                break;
            } else {
                cur = cur->right;
            }
        } else {
            // 暂时先不做任何动作
            free(new_node);
            return tree;
        }
    }
    tree->node_count++;

    // 添加完节点，调整平衡
    after_add_node(tree, new_node);
    
    return tree;
}



/**
 * @brief 删除叶子节点
 * 
 * @param tree 树
 * @param del 要删除的节点
 * @param flag 要删除的节点是其父节点的左节点还是右节点
 * @return int 
 */
int delete_leaf_node(avl_t* tree, avl_node_t* del, int flag)
{
    avl_node_t* del_parent = del->parent;

    /*! del == root，即要删除的叶子节点是根节点 */
    if (del == tree->root) {
        tree->root = NULL;
    } else {
        if (flag == LEFT_NODE) {
            del_parent->left = NULL;
        }
        if (flag == RIGHT_NODE) {
            del_parent->right = NULL;
        }
    }

    free(del);
    tree->node_count--;
    return 0;
}


/**
 * @brief Get the rsb min node object
 * 
 * @param root 
 * @param parentNode 
 * @return avl_node_t* 
 */
avl_node_t* get_rsb_min_node(avl_node_t* root, avl_node_t** parentNode)
{
    avl_node_t* cur;

    if (root == NULL) {
        *parentNode = NULL;
        return NULL;
    }

    *parentNode = root;
    cur = root->right;
    while (cur->left != NULL) {
        *parentNode = cur;
        cur = cur->left;
    }

    return cur;
}

/**
 * @brief 删除键值对应的节点
 * 
 * @param tree 
 * @param key 
 * @return int 
 */
int delete_node(avl_t* tree, int key)
{
    avl_node_t *del = NULL, *del_parent = NULL;
    avl_node_t *rsb_min_node, *rsb_min_node_parent;
    int flag;


    if (tree == NULL || tree->root == NULL) {
        return 0;
    }

    del = search_node(tree->root, key);
    if (del == NULL) {
        return 0;
    }
    /*! 根结点的父节点是空NULL */
    del_parent = del->parent;

    /*!< 叶子节点 */
    if (del->left == NULL && del->right == NULL) {
        /*! 是根节点 */
        if (del == tree->root) {
            delete_leaf_node(tree, del, -1);
            return 0;
        } else { /*! 不是父节点 */
            /*! 计算节点是其父节点的左节点还是右节点 */
            flag = key < del_parent->val ? LEFT_NODE : RIGHT_NODE;
            delete_leaf_node(tree, del, flag);

            // 删除节点后可能导致失衡
            after_delete_node(tree, del);
            return 0;
        }
    }
    else if (del->left == NULL && del->right != NULL) { /*!< 删除节点的左孩子为空，右孩子不为空 */
        if (del == tree->root) { /*! 根结点 */
            tree->root = tree->root->right;
            tree->root->parent = NULL;
        } else {
            /*! 计算节点是其父节点的左节点还是右节点 */
            flag = key < del_parent->val ? LEFT_NODE : RIGHT_NODE;
            if (flag == LEFT_NODE) {
                del->right->parent = del->parent;
                del_parent->left = del->right;        
            } else {
                del->right->parent = del->parent;
                del_parent->right = del->right;
            }
        }
        free(del);
        tree->node_count--;

        // 删除节点后可能导致失衡
        after_delete_node(tree, del);
    }
    else if (del->left != NULL && del->right == NULL) /*!< 删除节点的右孩子为空，左孩子不为空 */
    {
        if (del == tree->root) /*! 根结点 */
        {
            tree->root = tree->root->left;
            tree->root->parent = NULL;
        } else {
            /*! 计算节点是其父节点的左节点还是右节点 */
            flag = key < del_parent->val ? LEFT_NODE : RIGHT_NODE;
            if (flag == LEFT_NODE) {
                del->left->parent = del->parent;
                del_parent->left = del->left;
            } else {
                del->left->parent = del->parent;
                del_parent->right = del->left;
            }
        }
        free(del);
        tree->node_count--;

        // 删除节点后可能导致失衡
        after_delete_node(tree, del);
    } else { /*!< 删除节点的左右孩子都不为空 */
        rsb_min_node = get_rsb_min_node(del, &rsb_min_node_parent);
        if (del == tree->root) { /*! 根结点 */
            if (del->right->left == NULL) {
                rsb_min_node->left = tree->root->left;
                tree->root->left->parent = rsb_min_node;

                tree->root = rsb_min_node;
                tree->root->parent = NULL;
            } else {
                rsb_min_node_parent->left = rsb_min_node->right;
                if (rsb_min_node_parent->left == NULL) {
                    rsb_min_node->left = tree->root->left;
                    tree->root->left->parent = rsb_min_node;

                    rsb_min_node->right = tree->root->right;
                    tree->root->right->parent = rsb_min_node;

                    tree->root = rsb_min_node;
                    tree->root->parent = NULL;
                } else {
                    rsb_min_node_parent->left->parent = rsb_min_node_parent;

                    rsb_min_node->left = tree->root->left;
                    tree->root->left->parent = rsb_min_node;

                    rsb_min_node->right = tree->root->right;
                    tree->root->right->parent = rsb_min_node;

                    tree->root = rsb_min_node;
                    tree->root->parent = NULL;
                }
            }
            free(del);
            tree->node_count--;
            return 0;
        } else { /*!< 非根节点 */
            flag = key < del_parent->val ? LEFT_NODE : RIGHT_NODE;
            if (flag == LEFT_NODE) { /*! 要删除的节点是其父节点的左节点 */
                if (del == rsb_min_node_parent) {
                    rsb_min_node->parent = del->parent;
                    del->parent->left = rsb_min_node;

                    rsb_min_node->left = del->left;
                    del->left->parent = rsb_min_node;
                } else {
                    rsb_min_node_parent->left = rsb_min_node->right;
                    if (rsb_min_node_parent->left == NULL) {
                        rsb_min_node->left = del->left;
                        del->left->parent = rsb_min_node;

                        rsb_min_node->right = del->right;
                        del->right->parent = rsb_min_node;

                        del->parent->left = rsb_min_node;
                        rsb_min_node->parent = del->parent;
                    } else {
                        rsb_min_node_parent->left->parent = rsb_min_node_parent;

                        rsb_min_node->left = del->left;
                        del->left->parent = rsb_min_node;

                        rsb_min_node->right = del->right;
                        del->right->parent = rsb_min_node;

                        del->parent->left = rsb_min_node;
                        rsb_min_node->parent = del->parent;
                    }
                }
            }
            else { /*! 要删除的节点是其父节点的右节点 */
                if (del == rsb_min_node_parent) {
                    rsb_min_node->parent = del->parent;
                    del->parent->right = rsb_min_node;

                    rsb_min_node->left = del->left;
                    del->left->parent = rsb_min_node;
                } else {
                    rsb_min_node_parent->left = rsb_min_node->right;
                    if (rsb_min_node_parent->left == NULL) {
                        rsb_min_node->left = del->left;
                        del->left->parent = rsb_min_node;

                        rsb_min_node->right = del->right;
                        del->right->parent = rsb_min_node;

                        del->parent->right = rsb_min_node;
                        rsb_min_node->parent = del->parent;
                    } else {
                        rsb_min_node_parent->left->parent = rsb_min_node_parent;

                        rsb_min_node->left = del->left;
                        del->left->parent = rsb_min_node;

                        rsb_min_node->right = del->right;
                        del->right->parent = rsb_min_node;

                        del->parent->right = rsb_min_node;
                        rsb_min_node->parent = del->parent;
                    }
                }
            }
            free(del);
            tree->node_count--;

            // 删除节点后可能导致失衡
            after_delete_node(tree, del);
            return 0;
        }
    }

    return 0;
}


/**
 * @brief 计算树的高度
 * 
 * @param root 
 * @return int 
 */
int calc_tree_height(avl_node_t* root)
{
    if (root == NULL) {
        return 0;
    }

    int height = 0;

    avl_node_t* queue[10000];
    avl_node_t* tmp;
    int head = 0;
    int tail = 0;
    int count = 1;

    queue[tail++] = root;

    while (head != tail)
    {
        tmp = queue[head++];
        count--;

        if (tmp->left != NULL) {
            queue[tail++] = tmp->left;
        }
        if (tmp->right != NULL) {
            queue[tail++] = tmp->right;
        }

        if (count == 0) {
            height++;
            count = tail - head;
        }
    }

    return height;
}


/**
 * @brief 计算平衡因子
 * 
 * @param node 
 * @return int 
 */
int calc_balance_factor(avl_node_t* node)
{
    int lsb_height, rsb_height;
    
    // 传进来的节点为空，返回-1
    if (node == NULL) {
        return -1;
    }

    lsb_height = node->left == NULL ? 0 : node->left->height;
    rsb_height = node->right == NULL ? 0 : node->right->height;

    return lsb_height - rsb_height;
}


/**
 * @brief 节点是否平衡
 * 
 * @param node 
 * @return 1 ：平衡
 *         0 ：不平衡
 *         0：node为NULL
 */
int is_balance(avl_node_t* node)
{
    int res;

    if (node == NULL) {
        return 0;
    }

    // 计算节点的平衡因子
    res = calc_balance_factor(node);

    res = res < 0 ? (-res) : res;

    return res < 2 ? 1 : 0;
}


/**
 * @brief 计算两int型数较大值
 * 
 * @param a 
 * @param b 
 * @return int 
 */
int get_max_val(int a, int b)
{
    return a > b ? a : b;
}


/**
 * @brief 更新节点高度
 * 
 * @param node 
 * @return void
 */
void updata_node_height(avl_node_t* node)
{
    int left_height, right_height;

    if (node == NULL) {
        return;
    }

    left_height = node->left == NULL ? 0 : node->left->height;
    right_height = node->right == NULL ? 0 : node->right->height;

    node->height = 1 + get_max_val(left_height, right_height);
}


/**
 * @brief 是否是左孩子
 * 
 * @param node 
 * @return int 
 */
int is_left_child(avl_node_t* node)
{
    if(node == NULL) {
        return 0;
    }

    return node->parent != NULL && node == node->parent->left;
}


/**
 * @brief 是否是右孩子
 * 
 * @param node 
 * @return int 
 */
int is_right_child(avl_node_t* node)
{
    if (node == NULL) {
        return 0;
    }

    return node->parent != NULL && node == node->parent->right;
}


/**
 * @brief 找出当前节点高度更高的孩子节点
 * 
 * @param tree 
 * @param node 
 * @return avl_node_t* 
 */
avl_node_t* taller_child(avl_t* tree, avl_node_t* node)
{
    int lsb_height, rsb_height;

    if(tree == NULL || node == NULL) {
        return NULL;
    }

    lsb_height = node->left == NULL ? 0 : node->left->height;
    rsb_height = node->right == NULL ? 0 : node->right->height;

    if (lsb_height > rsb_height) {
        return node->left;
    } else if(lsb_height < rsb_height)
    {
        return node->right;
    } else
    {
        return is_left_child(node) ? node->left : node->right;
    }
}


/**
 * @brief 左旋右旋后半部分步骤相同，单独封装函数
 * 
 * @param node 
 */
static void after_rotate(avl_t* tree, avl_node_t* grandpa, avl_node_t* parent, avl_node_t* child)
{
    parent->parent = grandpa->parent;
    if (is_left_child(grandpa)) {
        grandpa->parent->left = parent;
    } else if (is_right_child(grandpa)) {
        grandpa->parent->right = parent;
    } else {
        tree->root = parent;
        tree->root->parent = NULL;
    }

    if (child != NULL) {
        child->parent = grandpa;
    }

    grandpa->parent = parent;

    updata_node_height(grandpa);
    updata_node_height(parent);
}


/**
 * @brief 将此节点左旋
 * 
 * @param grandpa_node 
 * @param parent 
 * @return int 
 */
int rotate_left(avl_t* tree, avl_node_t* grandpa)
{
    avl_node_t *parent, *child;

    if (grandpa == NULL) {
        return -1;
    }

    parent = grandpa->right;
    if (parent == NULL) {
        return -1;
    }

    child = parent->left;
    grandpa->right = child;
    parent->left = grandpa;

    after_rotate(tree, grandpa, parent, child);

    return 0;
}


/**
 * @brief 将此节点右旋
 * 
 * @param grandpa_node 
 * @param parent 
 * @return int 
 */
int rotate_right(avl_t* tree, avl_node_t* grandpa)
{
    avl_node_t *parent, *child;

    if (grandpa == NULL) {
        return -1;
    }


    parent = grandpa->left;
    if (parent == NULL) {
        return -1;
    }

    child = parent->right;
    grandpa->left = child;
    parent->right = grandpa;

    after_rotate(tree, grandpa, parent, child);

    return 0;
}


/**
 * @brief 节点失衡后，恢复节点的平衡
 * 
 * @param tree 
 * @param node 高度最低的不平衡的节点
 * @return int 
 */
int rebalance_node(avl_t* tree, avl_node_t* grandpa)
{
    avl_node_t *parent, *node;

    if (tree == NULL || grandpa == NULL) {
        return -1;
    }

    parent = taller_child(tree, grandpa);
    node = taller_child(tree, parent);

    //parent = 

    if (is_left_child(parent)) {
        // LL右旋
        if (is_left_child(node)) {
            rotate_right(tree, grandpa);
        } else {
            // LR先左旋再右旋
            rotate_left(tree, parent);
            rotate_right(tree, grandpa);
        }
    } else if (is_right_child(parent)) {
        // RL先右旋再左旋
        if (is_left_child(node)) {
            rotate_right(tree, parent);
            rotate_left(tree, grandpa);
        } else {
            // RR左旋
            rotate_left(tree, grandpa);
        }
    }

    return 0;
}



int maxDepth(avl_node_t* root)
{
    if (root == NULL) {
        return 0;
    }

    return get_max_val(maxDepth(root->left), maxDepth(root->right)) + 1;
}


/**
 * @brief 处理添加节点后导致失衡
 * 
 * @param tree 
 * @param node 新添加的节点
 * @return int 
 */
int after_add_node(avl_t* tree, avl_node_t* node)
{
    avl_node_t* cur = node;

    if (tree == NULL || node == NULL) {
        return -1;
    }

    // 添加的是根结点，不会导致失衡，直接退出
    if (node == tree->root) {
        return 0;
    }

    while ((cur = cur->parent) != NULL) {
        // 如果此节点平衡，则更新高度，否则将节点恢复平衡
        if (is_balance(cur)) {
            updata_node_height(cur);
        } else {
            rebalance_node(tree, cur);
            break;
        }
    }

    return 0;
}


/**
 * @brief 处理删除节点后导致失衡
 * 
 * @param tree 
 * @param node 要删除的节点
 * @return int 
 */
int after_delete_node(avl_t* tree, avl_node_t* node)
{
    avl_node_t* cur = node;

    if (tree == NULL || node == NULL) {
        return -1;
    }

    while ((cur = cur->parent) != NULL) {
        // 如果此节点平衡，则更新高度，否则将节点恢复平衡
        if (is_balance(cur)) {
            updata_node_height(cur);
        } else {
            rebalance_node(tree, cur);
        }
    }

    return 0;
}



void infix(avl_node_t* root)
{
    if (root == NULL) {
        return;
    }

    infix(root->left);
    printf("%d ", root->val);
    infix(root->right);
}

int level(avl_node_t* root)
{
    if (root == NULL) {
        return 0;
    }

    avl_node_t* queue[10000];
    avl_node_t* tmp;
    int head = 0;
    int tail = 0;

    queue[tail++] = root;

    while (head != tail) {
        tmp = queue[head++];
        printf("%d ", tmp->val);

        if (tmp->left != NULL) {
            queue[tail++] = tmp->left;
        }

        if (tmp->right != NULL) {
            queue[tail++] = tmp->right;
        }
    }

    return 0;
}

